import komand
from .schema import SearchForExploitInput, SearchForExploitOutput
# Custom imports below


class SearchForExploit(komand.Action):

    def __init__(self):
        super(self.__class__, self).__init__(
                name='search_for_exploit',
                description='Searches for an exploit with a Metasploit RPC session',
                input=SearchForExploitInput(),
                output=SearchForExploitOutput())

    def run(self, params={}):
        search_term = params["search_term"].lower().encode("utf-8", "ignore")
        results = {"exploits": [], "auxiliaries": [], "post_exploits": []}
        client = self.connection.client
        try:
            exploits = client.modules.exploits
            auxiliary = client.modules.auxiliary
            post_exploits = client.modules.post

            for exploit in exploits:
                current_exploit = client.modules.use("exploit", exploit)
                description = current_exploit.description
                if search_term.lower() in exploit or search_term.lower() in description.lower():
                    results["exploits"].append(exploit)

            for aux in auxiliary:
                current_aux = client.modules.use("auxiliary", aux)
                description = current_aux.description
                try:
                    if search_term in aux or search_term in description.lower():
                        results["auxiliaries"].append(aux)
                except UnicodeDecodeError as e:
                    self.logger.error("Error: %s" % e)
            for post in post_exploits:
                current_post = client.modules.use("post", post)
                description = current_post.description
                if search_term.lower() in post or search_term.lower() in description.lower():
                    results["post_exploits"].append(post)
            return {"search_results": results}
        except Exception as e:
            self.logger.error("Error occurred while retrieving exploits: ", e)
            raise

    def test(self):
        client = self.connection.client
        try:
            exploits = client.modules.exploits
        except Exception as e:
            self.logger.error("Error occurred while testing search for exploits: ", e)
            raise
        return {"success": True}
